// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <memory>
#include <vector>
#include <pollux/dwio/common/scan_spec.h>
#include <pollux/type/type.h>

namespace kumo::pollux::dwio::common {
    class TypeWithId : public pollux::Tree<std::shared_ptr<const TypeWithId> > {
    public:
        /// NOTE: This constructor will re-parent the children.
        TypeWithId(
            std::shared_ptr<const pollux::Type> type,
            std::vector<std::unique_ptr<TypeWithId> > &&children,
            uint32_t id,
            uint32_t maxId,
            uint32_t column);

        TypeWithId(const TypeWithId &) = delete;

        TypeWithId &operator=(const TypeWithId &) = delete;

        static std::unique_ptr<TypeWithId> create(
            const std::shared_ptr<const pollux::Type> &root,
            uint32_t next = 0);

        /// Create TypeWithId node but leave all the unselected children as nullptr.
        /// The ids are set correctly even when some of the previous nodes are not
        /// selected.
        static std::unique_ptr<TypeWithId> create(
            const RowTypePtr &type,
            const pollux::common::ScanSpec &spec);

        uint32_t size() const override;

        const std::shared_ptr<const pollux::Type> &type() const {
            return type_;
        }

        const TypeWithId *parent() const {
            return parent_;
        }

        uint32_t id() const {
            return id_;
        }

        uint32_t maxId() const {
            return maxId_;
        }

        uint32_t column() const {
            return column_;
        }

        const std::shared_ptr<const TypeWithId> &childAt(uint32_t idx) const override;

        const std::shared_ptr<const TypeWithId> &childByName(
            const std::string &name) const {
            POLLUX_CHECK_EQ(type_->kind(), pollux::TypeKind::ROW);
            return childAt(type_->as<pollux::TypeKind::ROW>().getChildIdx(name));
        }

        const std::vector<std::shared_ptr<const TypeWithId> > &getChildren() const {
            return children_;
        }

        std::string fullName() const;

    private:
        static std::unique_ptr<TypeWithId> create(
            const std::shared_ptr<const pollux::Type> &type,
            uint32_t &next,
            uint32_t column);

        const std::shared_ptr<const pollux::Type> type_;
        const TypeWithId *const parent_;
        const uint32_t id_;
        const uint32_t maxId_;
        const uint32_t column_;
        const std::vector<std::shared_ptr<const TypeWithId> > children_;
    };
} // namespace kumo::pollux::dwio::common
